"""
任务四：迁移诊断可解释性分析主程序

基于任务三的DANN域适应模型，结合轴承故障特点与故障机理，
对迁移诊断的可解释性进行全面分析。

作者：数学建模团队
版本：1.0
"""

import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# 添加父目录到路径
sys.path.append('..')

# 导入任务四模块
from pytorch_interpretability_analysis import (
    load_task3_data, create_dann_model, InterpretabilityAnalyzer, 
    InterpretabilityVisualizer, generate_interpretability_report
)
from task4_config import *

# 设置matplotlib非交互式后端
plt.switch_backend('Agg')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


def main():
    """主函数"""
    print("🚀 开始执行任务四：迁移诊断可解释性分析")
    print("=" * 80)
    
    # 生成时间戳
    timestamp = datetime.now().strftime(OUTPUT_CONFIG['timestamp_format'])
    print(f"⏰ 执行时间: {timestamp}")
    
    try:
        # 1. 加载任务三数据
        print("\n" + "=" * 60)
        print("步骤1: 加载任务三数据")
        print("=" * 60)
        
        X_source, y_source, X_target, feature_names = load_task3_data()
        
        # 2. 创建DANN模型
        print("\n" + "=" * 60)
        print("步骤2: 创建DANN模型")
        print("=" * 60)
        
        input_dim = X_source.shape[1]
        num_classes = len(np.unique(y_source))
        model = create_dann_model(input_dim, num_classes)
        
        # 3. 创建可解释性分析器
        print("\n" + "=" * 60)
        print("步骤3: 创建可解释性分析器")
        print("=" * 60)
        
        analyzer = InterpretabilityAnalyzer(model, MODEL_CONFIG['device'])
        visualizer = InterpretabilityVisualizer()
        
        # 4. 事前可解释性分析
        print("\n" + "=" * 60)
        print("步骤4: 事前可解释性分析")
        print("=" * 60)
        
        print("🔍 分析特征重要性...")
        feature_importance = analyzer.analyze_feature_importance(X_source, y_source, feature_names)
        
        print("🔍 分析故障特征...")
        fault_analysis = analyzer.analyze_fault_characteristics(X_source, y_source, feature_names)
        
        # 5. 迁移过程可解释性分析
        print("\n" + "=" * 60)
        print("步骤5: 迁移过程可解释性分析")
        print("=" * 60)
        
        print("🔍 分析域适应过程...")
        domain_analysis = analyzer.analyze_domain_adaptation(X_source, X_target)
        
        print("🔍 分析层激活...")
        layer_activations = analyzer.analyze_layer_activations(X_source[:1])
        
        # 6. 事后可解释性分析
        print("\n" + "=" * 60)
        print("步骤6: 事后可解释性分析")
        print("=" * 60)
        
        print("🔍 分析梯度重要性...")
        gradient_importance = analyzer.analyze_gradient_importance(X_source[:1], y_source[0])
        
        print("🔍 分析决策边界...")
        decision_boundary = analyzer.analyze_decision_boundary(X_source, y_source, X_target, feature_names)
        
        # 7. 生成可视化图表
        print("\n" + "=" * 60)
        print("步骤7: 生成可视化图表")
        print("=" * 60)
        
        print("📊 生成特征重要性分析图...")
        visualizer.plot_feature_importance(
            feature_importance, 
            top_n=INTERPRETABILITY_CONFIG['feature_importance']['top_n'],
            save_path=f'feature_importance_analysis_{timestamp}.png'
        )
        
        print("📊 生成域适应过程分析图...")
        visualizer.plot_domain_adaptation_analysis(
            domain_analysis, 
            save_path=f'domain_adaptation_analysis_{timestamp}.png'
        )
        
        print("📊 生成故障特征分析图...")
        visualizer.plot_fault_characteristics_analysis(
            fault_analysis, 
            feature_names, 
            save_path=f'fault_characteristics_analysis_{timestamp}.png'
        )
        
        print("📊 生成决策边界分析图...")
        visualizer.plot_decision_boundary_analysis(
            decision_boundary, 
            save_path=f'decision_boundary_analysis_{timestamp}.png'
        )
        
        # 8. 生成综合报告
        print("\n" + "=" * 60)
        print("步骤8: 生成综合报告")
        print("=" * 60)
        
        report_path = generate_interpretability_report(
            analyzer, visualizer, X_source, y_source, X_target, feature_names, timestamp
        )
        
        # 9. 保存分析数据
        print("\n" + "=" * 60)
        print("步骤9: 保存分析数据")
        print("=" * 60)
        
        # 保存特征重要性数据
        feature_importance.to_csv(f'feature_importance_data_{timestamp}.csv', index=False, encoding='utf-8-sig')
        print("✅ 特征重要性数据已保存")
        
        # 保存域适应分析数据
        domain_data = {
            'euclidean_distance': domain_analysis['euclidean_distance'],
            'cosine_distance': domain_analysis['cosine_distance'],
            'js_divergence': domain_analysis['js_divergence']
        }
        pd.DataFrame([domain_data]).to_csv(f'domain_adaptation_data_{timestamp}.csv', index=False, encoding='utf-8-sig')
        print("✅ 域适应分析数据已保存")
        
        # 保存故障特征分析数据
        fault_data = []
        for fault_type, stats in fault_analysis.items():
            fault_data.append({
                'fault_type': fault_type,
                'mean_mean': np.mean(stats['mean']),
                'mean_std': np.mean(stats['std']),
                'mean_max': np.mean(stats['max']),
                'mean_min': np.mean(stats['min'])
            })
        pd.DataFrame(fault_data).to_csv(f'fault_characteristics_data_{timestamp}.csv', index=False, encoding='utf-8-sig')
        print("✅ 故障特征分析数据已保存")
        
        print("\n" + "=" * 80)
        print("🎉 任务四执行完成！")
        print("=" * 80)
        print("📁 生成的文件:")
        print(f"  📊 特征重要性分析图: feature_importance_analysis_{timestamp}.png")
        print(f"  📊 域适应过程分析图: domain_adaptation_analysis_{timestamp}.png")
        print(f"  📊 故障特征分析图: fault_characteristics_analysis_{timestamp}.png")
        print(f"  📊 决策边界分析图: decision_boundary_analysis_{timestamp}.png")
        print(f"  📝 可解释性分析报告: {report_path}")
        print(f"  📄 特征重要性数据: feature_importance_data_{timestamp}.csv")
        print(f"  📄 域适应分析数据: domain_adaptation_data_{timestamp}.csv")
        print(f"  📄 故障特征分析数据: fault_characteristics_data_{timestamp}.csv")
        print("=" * 80)
        
        # 10. 输出关键发现
        print("\n" + "=" * 60)
        print("关键发现总结")
        print("=" * 60)
        
        # 特征重要性发现
        top_feature = feature_importance.iloc[0]
        print(f"🔍 最重要特征: {top_feature['feature']} (重要性: {top_feature['importance']:.4f})")
        
        # 域适应效果
        adaptation_score = 1 - domain_analysis['cosine_distance']
        print(f"🔍 域适应效果: {adaptation_score:.4f} (1为完美适应)")
        
        # 故障特征发现
        fault_types = list(fault_analysis.keys())
        print(f"🔍 分析的故障类型: {', '.join(fault_types)}")
        
        print("\n✅ 任务四可解释性分析完成！")
        
    except Exception as e:
        print(f"❌ 执行过程中出现错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

